#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <gflags/gflags.h>
#include <glog/logging.h>

#include "core/grape.h"
#include "CCTemporal.hpp"

using namespace std;

template<class VD, class ED, class MessageT, class ResultT, class TemporalT>
class CCWorker: public Worker<VD, ED, MessageT, ResultT, TemporalT> {

	void pEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {
		unordered_map<VertexID, VertexID> localCompID;
		unordered_map<VertexID, VertexID> globalCompID;
		unordered_map<VertexID, vector<VertexID>> outerNodes;
		unordered_map<VertexID, VertexID> messages;

//		vector<VertexID> innerVertices = fragment.getInnerVerticesID();
//		for (auto x : innerVertices) {
		for (auto x : *fragment.getInnerVertices()) {
			if (localCompID.find(x) != localCompID.end()) {
				continue;
			}
			queue<VertexID> que;
			que.push(x);
			localCompID[x] = INT_MAX;
			queue<VertexID> comp;
			vector<VertexID> outers;
			VertexID minID = x;
			for (; !que.empty(); que.pop()) {
				VertexID u = que.front();
				comp.push(u);

				for (auto v : fragment.getChildrenID(u)) {
					minID = min(minID, v);
					if (fragment.isInnerVertex(v)) {
						if (localCompID.find(v) == localCompID.end()) {
							que.push(v);
							localCompID[v] = INT_MAX;
						}
					} else {
						outers.push_back(v);
					}
				}

				for (auto v : fragment.getParentsID(u)) {
					minID = min(minID, v);
					if (fragment.isInnerVertex(v)) {
						if (localCompID.find(v) == localCompID.end()) {
							que.push(v);
							localCompID[v] = INT_MAX;
						}
					} else {
						outers.push_back(v);
					}
				}
			}
			for (; !comp.empty(); comp.pop()) {
				VertexID u = comp.front();
				localCompID[u] = minID;
			}
			globalCompID[minID] = minID;
			outerNodes[minID] = outers;
			for (auto v : outers) {
				if (messages.find(v) == messages.end() || minID < messages[v]) {
					messages[v] = minID;
				}
			}
		}

		for (auto p : messages) {
			VertexID u = p.first;
			VertexID compID = p.second;
			int fID = fragment.getFragmentID(u);
			messageBuffer.add_message(fID, PairMessage<VertexID>(u, compID));
		}

		temporal = CCTemporal(localCompID, globalCompID, outerNodes);

		partialResult = make_pair(localCompID, globalCompID);
	}

	void incEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {
		auto localCompID = temporal.localCompID;
		auto globalCompID = temporal.globalCompID;
		auto outerNodes = temporal.outerNodes;

		unordered_set<VertexID> updated;
		vector<MessageT> received = messageBuffer.get_messages();
		for (auto item : received) {
			VertexID u = item.getID();
			VertexID tag = item.getValue();
			VertexID compID = localCompID[u];
			if (globalCompID[compID] > tag) {
				globalCompID[compID] = tag;
				if (updated.find(compID) == updated.end()) {
					updated.insert(compID);
				}
			}
		}

		unordered_map<VertexID, VertexID> messages;
		for (auto u : updated) {
			VertexID tag = globalCompID[u];
			for (auto v : outerNodes[u]) {
				if (messages.find(v) == messages.end() || messages[v] > tag) {
					messages[v] = tag;
				}
			}
		}

		for (auto p : messages) {
			VertexID u = p.first;
			VertexID compID = p.second;
			int fID = fragment.getFragmentID(u);
			messageBuffer.add_message(fID, PairMessage<VertexID>(u, compID));
		}

		temporal.globalCompID = globalCompID;
		partialResult.second = globalCompID;
	}

	void assemble(vector<ResultT> &partialResults) {
		if (this->result_file != "") {
			std::ofstream fout;
			fout.open(this->result_file);
			for (auto res : partialResults) {
				auto localCompID = res.first;
				auto globalCompID = res.second;
				for (auto p : localCompID) {
					VertexID u = p.first;
					VertexID compID = p.second;
					VertexID minHash = globalCompID[compID];
					fout << u << '\t' << minHash << "\n";
				}
			}
			fout.close();
		}
	}
};

int main(int argc, char *argv[]) {

	FLAGS_stderrthreshold = 0;

	google::SetUsageMessage("Usage: mpiexec.hydra [mpi_opts] cc [grape_opts]");
	google::ParseCommandLineFlags(&argc, &argv, true);
	init_workers();

	WorkerParams param;

	param.algorithm = "cc";
	param.use_pattern = false;
	param.enable_assemble = true;
	param.graph_load_strategy = BOTH_OUT_IN;

	google::ShutDownCommandLineFlags();

	CCWorker<VertexID, int, PairMessage<VertexID>, pair<unordered_map<VertexID, VertexID>, unordered_map<VertexID, VertexID>>, CCTemporal> ccWorker;
	google::InitGoogleLogging("cc");
	ccWorker.run(param);
	google::ShutdownGoogleLogging();
	worker_finalize();
	return 0;
}
